import torch.nn as nn

class AutoEncoder(nn.Module):
	def __init__(self, in_sz, hd_sz):
		super(AutoEncoder, self).__init__()

		base = hd_sz
		self.encoder = nn.Sequential(
			nn.Linear(in_sz, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, base),
			nn.ReLU(),
			)
		self.decoder = nn.Sequential(
			nn.BatchNorm1d(base),
			nn.Linear(base, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, in_sz),
			# nn.ReLU()
			)

	def forward(self, x):
		z = self.encoder(x)
		return self.decoder(z)
	
			